# import torch

# batch_size = 1
# seq_len = 3
# input_size = 4
# hidden_size = 2
#
# cell = torch.nn.RNNCell(input_size=input_size, hidden_size=hidden_size)
#
# dataset = torch.randn(seq_len, batch_size, input_size)
# hidden = torch.zeros(batch_size, hidden_size)
# for idx, input in enumerate(dataset):
#     print('='*20, idx, '='*20)
#     print('Input size:', input.shape)
#     hidden = cell(input, hidden)
#     print('outputs size:', hidden.shape)
#     print(hidden)

# _______________________________________________________________________________________________
# import torch
# batch_size = 1
# seq_len = 3
# input_size = 4
# hidden_size = 2
# num_layers = 1
# cell = torch.nn.RNN(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
# inputs = torch.randn(seq_len, batch_size, input_size)
# hidden = torch.zeros(num_layers, batch_size, hidden_size)
# out, hidden = cell(inputs, hidden)
# print('Output size:', out.shape)
# print('Output:', out)
# print("hidden size:", hidden.shape)
# print('Hidden:', hidden)


# -------------------------------------------------------------------------------------
import torch
import torch.nn as nn

input_size = 4
hidden_size = 4
batch_size = 1
idx2char = ['e', 'h', 'l', 'o']
x_data = [1, 0, 2, 2, 3]
y_data = [3, 1, 2, 3, 2]
one_hot_lookup = [[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 0, 1],
                  [0, 0, 0, 1]]
x_one_hot = [one_hot_lookup[x] for x in x_data]
inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)
labels = torch.LongTensor(y_data).view(-1, 1)

class Model(nn.Module):
    def __init__(self, input_size, hidden_size, batch_size):
        super(Model, self).__init__()
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnncell = nn.RNNCell(input_size=input_size, hidden_size=hidden_size)

    def forward(self, input, hidden):
        hidden = self.rnncell(input, hidden)
        return hidden

    def init_hidden(self):
        return torch.zeros(self.batch_size, self.hidden_size)

net = Model(input_size, hidden_size, batch_size)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.1)

for epoch in range(50):
    loss = 0
    optimizer.zero_grad()
    hidden = net.init_hidden()
    print('Predicted string:', end='')
    for input, label in zip(inputs, labels):
        hidden = net(input, hidden)
        loss += criterion(hidden, label)
        _, idx = hidden.max(dim=1)
        print(idx2char[idx.item()], end='')
    loss.backward()
    optimizer.step()
    print(', Epoch [%d/50] loss=%.4f' % (epoch+1, loss.item()))
